#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@his-pc>
#
# Distributed under terms of the MIT license.

"""

"""

class Connection:
    def __init__(self, city1, city2, cost):
        self.city1, self.city2, self.cost = city1, city2, cost

    def __str__(self):
        return self.city1 + " " + self.city2 + " " + str(self.cost)


class GraphNode:
    def __init__(self, city_name=None):
        self.city = city_name
        self.neighbors = {}     # { city_name : (weight, GraphNode)}
        self.dis = 12345567890  # 距离visited树中距离
        self.dis_conn = ''  # 距离visited树中距离

    def add(self, conn, node, weight):
        # self.neighbors[node.city] = (weight, node)
        self.neighbors[node.city] = (conn, node, weight)


class Solution:
    # @param {Connection[]} connections given a list of connections
    # include two cities and cost
    # @return {Connection[]} a list of connections from results
    def lowestCost(self, connections):
        # return self.kruskal_lowest_cost(connections)
        return self.prim_lowest_cost(connections)

    # =================== prim method =====================
    def prim_lowest_cost(self, connections):
        self.connections = connections
        self.non_visited = {}     # {city_name : node}, adjacent list
        self.build_adjacency_list()

        first_name, first_node = self.non_visited.popitem()
        self.update_dis(first_node)
        self.result = []
        # print(len(self.non_visited))

        while self.non_visited:
            min_node = self.find_min_node()
            if not min_node:
                # print('not found')
                return []

            # print('found %s' % min_node.city)
            self.result.append(min_node.dis_conn)
            self.update_dis(min_node)

        self.result.sort(key=lambda c:(c.cost, c.city1, c.city2))
        return self.result

    def find_min_node(self):
        min_dis = 1234567890 + 1
        min_node = None
        # 这里还要加入城市名称比较, 符合题意

        for city_name, node in self.non_visited.items():
            # print('in find %s %s' % (node.city, node.dis))
            if node.dis < min_dis:
                min_node = node
                min_dis = node.dis
                continue

            if node.dis == min_dis:
                print('equal %s %s %s %s %s' % (node.dis, node.dis_conn.city1, node.dis_conn.city2, min_node.dis_conn.city1, min_node.dis_conn.city2))
                if (node.dis_conn.city1, node.dis_conn.city2) < (min_node.dis_conn.city1, min_node.dis_conn.city2):
                    min_node = node

        if min_node:
            # print('find %s %s' % (node.city, node.dis))
            self.non_visited.pop(min_node.city)

        return min_node

    def update_dis(self, selected_node):
        for city_name, (conn, node, weight) in selected_node.neighbors.items():
            if city_name not in self.non_visited:
                continue

            if node.dis > weight:
                # print('update %s %s' % (node.city, weight))
                node.dis = weight
                node.dis_conn = conn
        # print('update ====')

    def build_adjacency_list(self):
        for conn in self.connections:
            start, end, weight = conn.city1, conn.city2, conn.cost
            start_node = self.get_city(start)
            end_node   = self.get_city(end)
            start_node.add(conn, end_node, weight)
            end_node.add(conn, start_node, weight)

    def get_city(self, name):
        city = self.non_visited.get(name)
        if city is None:
            city = GraphNode(name)
            self.non_visited[name] = city
        return city

    # =================== kruskal method =====================
    def kruskal_lowest_cost(self, connections):
        # 因为heapq不支持key参数, 所以无法对Connection类基于cost 进行heapify()
        # 这里使用排序替代
        self.cities = set()
        for conn in connections:
            start, end, cost = conn.city1, conn.city2, conn.cost
            self.cities.add(start)
            self.cities.add(end)

        self.vertex_num = len(self.cities)

        # 题目要求在cost相同时,按照城市排序
        connections.sort(key=lambda c: (c.cost, c.city1, c.city2), reverse=True)

        # self.disjoint_parent = {i:-1 for i in range(1, vertex_num + 1)}        # make union_ set
        self.disjoint_parent = {}        # make union_set
        self.tree_edge_cnt = 0                                                  # 也可以换成已经在树中节点个数
        self.minimal_edges = []

        while not self.span_over and connections:
            # weight, start, end = cur_edge = heapq.heappop(self.edges)
            conn = connections.pop()        # 从后面弹出, 效率会高一些
            start, end, cost = conn.city1, conn.city2, conn.cost
            if not self.union_(start, end):
                continue

            self.tree_edge_cnt += 1
            self.minimal_edges.append(conn)

        if self.span_over:
            return self.minimal_edges
        else:
            return []

    @property
    def span_over(self):
        # 是否匹配结束
        return self.tree_edge_cnt == self.vertex_num - 1

    def find_parent(self, node_label):
        # 注意: 这里递归压缩parent路径
        # parent_label = self.disjoint_parent[node_label]
        parent_label = self.disjoint_parent.get(node_label)

        if parent_label is None:
            self.disjoint_parent[node_label] = node_label
            return self.disjoint_parent[node_label]

        if parent_label == node_label:
            return node_label

        self.disjoint_parent[node_label] = self.find_parent(parent_label)
        return self.disjoint_parent[node_label]

    def union_(self, n1, n2):
        p1 = self.find_parent(n1)
        p2 = self.find_parent(n2)

        if p1 == p2:
            return False

        self.disjoint_parent[p1] = p2
        return True


def main():
    print("start main")
    s = Solution()
    ret = s.lowestCost([
        # Connection('a', 'b', 1),
        # Connection('b', 'c', 2),
        # Connection('c', 'a', 3),
        Connection('a', 'b', 1),
        Connection('a', 'c', 2),
        Connection('b', 'c', 3),
    ])

    print(ret[1])
    print(ret[0])

if __name__ == "__main__":
    main()
